import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np
import time

# 设置随机种子，确保实验可重复
torch.manual_seed(42)
np.random.seed(42)

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # 标准化处理
])

# 加载FashionMNIST数据集
fashion_train = datasets.FashionMNIST('./data', train=True, download=True, transform=transform)
fashion_test = datasets.FashionMNIST('./data', train=False, transform=transform)

# 加载CIFAR-10数据集
cifar_train = datasets.CIFAR10('./data', train=True, download=True, transform=transform)
cifar_test = datasets.CIFAR10('./data', train=False, transform=transform)

# 创建数据加载器
batch_size = 64
fashion_train_loader = DataLoader(fashion_train, batch_size=batch_size, shuffle=True)
fashion_test_loader = DataLoader(fashion_test, batch_size=batch_size)
cifar_train_loader = DataLoader(cifar_train, batch_size=batch_size, shuffle=True)
cifar_test_loader = DataLoader(cifar_test, batch_size=batch_size)

# 定义3层神经网络模型
class ThreeLayerNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes, dataset_type):
        super(ThreeLayerNet, self).__init__()
        self.dataset_type = dataset_type
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, num_classes)
        
    def forward(self, x):
        # 根据数据集类型调整输入形状
        if self.dataset_type == 'fashion':
            x = x.view(-1, 28 * 28)  # FashionMNIST图像尺寸为28x28
        else:
            x = x.view(-1, 3 * 32 * 32)  # CIFAR-10图像尺寸为32x32，3通道
            
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.fc3(out)
        return out

# 训练函数
def train_model(model, train_loader, test_loader, criterion, optimizer, epochs, device):
    train_losses = []
    test_losses = []
    train_accuracies = []
    test_accuracies = []
    
    for epoch in range(epochs):
        # 训练阶段
        model.train()
        train_loss = 0
        correct = 0
        total = 0
        
        for inputs, targets in train_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            
            train_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
        
        train_loss /= len(train_loader)
        train_accuracy = 100. * correct / total
        train_losses.append(train_loss)
        train_accuracies.append(train_accuracy)
        
        # 测试阶段
        model.eval()
        test_loss = 0
        correct = 0
        total = 0
        
        with torch.no_grad():
            for inputs, targets in test_loader:
                inputs, targets = inputs.to(device), targets.to(device)
                
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                
                test_loss += loss.item()
                _, predicted = outputs.max(1)
                total += targets.size(0)
                correct += predicted.eq(targets).sum().item()
        
        test_loss /= len(test_loader)
        test_accuracy = 100. * correct / total
        test_losses.append(test_loss)
        test_accuracies.append(test_accuracy)
        
        print(f'Epoch {epoch+1}/{epochs} | '
              f'Train Loss: {train_loss:.4f} | Train Acc: {train_accuracy:.2f}% | '
              f'Test Loss: {test_loss:.4f} | Test Acc: {test_accuracy:.2f}%')
    
    return train_losses, test_losses, train_accuracies, test_accuracies

# 可视化函数
def visualize_results(train_losses, test_losses, train_accuracies, test_accuracies, dataset_name):
    plt.figure(figsize=(12, 4))
    
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Train Loss')
    plt.plot(test_losses, label='Test Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title(f'{dataset_name} - Loss Curve')
    plt.legend()
    plt.grid(True)
    
    plt.subplot(1, 2, 2)
    plt.plot(train_accuracies, label='Train Accuracy')
    plt.plot(test_accuracies, label='Test Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title(f'{dataset_name} - Accuracy Curve')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    plt.savefig(f'{dataset_name}_metrics.png')
    plt.show()

# 可视化预测结果
def visualize_predictions(model, test_loader, classes, dataset_type, device, num_samples=16):
    model.eval()
    images, labels = next(iter(test_loader))
    images, labels = images[:num_samples].to(device), labels[:num_samples].to(device)
    
    with torch.no_grad():
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)
    
    plt.figure(figsize=(10, 10))
    for i in range(num_samples):
        plt.subplot(4, 4, i+1)
        if dataset_type == 'fashion':
            img = images[i].cpu().view(28, 28)
        else:
            img = images[i].cpu().permute(1, 2, 0)
        
        # 反标准化处理
        img = img * 0.5 + 0.5
        
        plt.imshow(img, cmap='gray' if dataset_type == 'fashion' else None)
        plt.title(f'Pred: {classes[predicted[i]]}\nTrue: {classes[labels[i]]}', fontsize=9)
        plt.axis('off')
    
    plt.tight_layout()
    plt.savefig(f'{dataset_type}_predictions.png')
    plt.show()

# 设置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'Using device: {device}')

# 训练FashionMNIST模型
fashion_model = ThreeLayerNet(input_size=28*28, hidden_size=128, num_classes=10, dataset_type='fashion').to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(fashion_model.parameters(), lr=0.001)

print("\nTraining FashionMNIST Model...")
start_time = time.time()
fashion_train_losses, fashion_test_losses, fashion_train_acc, fashion_test_acc = train_model(
    fashion_model, fashion_train_loader, fashion_test_loader, criterion, optimizer, epochs=20, device=device
)
fashion_training_time = time.time() - start_time
print(f"FashionMNIST training completed in {fashion_training_time:.2f} seconds")

# 训练CIFAR-10模型
cifar_model = ThreeLayerNet(input_size=3*32*32, hidden_size=128, num_classes=10, dataset_type='cifar').to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(cifar_model.parameters(), lr=0.001)

print("\nTraining CIFAR-10 Model...")
start_time = time.time()
cifar_train_losses, cifar_test_losses, cifar_train_acc, cifar_test_acc = train_model(
    cifar_model, cifar_train_loader, cifar_test_loader, criterion, optimizer, epochs=20, device=device
)
cifar_training_time = time.time() - start_time
print(f"CIFAR-10 training completed in {cifar_training_time:.2f} seconds")

# 保存模型
torch.save(fashion_model.state_dict(), 'fashion_model.pth')
torch.save(cifar_model.state_dict(), 'cifar_model.pth')

# 可视化结果
fashion_classes = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 
                   'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
cifar_classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                 'dog', 'frog', 'horse', 'ship', 'truck']

visualize_results(fashion_train_losses, fashion_test_losses, fashion_train_acc, fashion_test_acc, 'FashionMNIST')
visualize_results(cifar_train_losses, cifar_test_losses, cifar_train_acc, cifar_test_acc, 'CIFAR-10')

# 可视化预测结果
visualize_predictions(fashion_model, fashion_test_loader, fashion_classes, 'fashion', device)
visualize_predictions(cifar_model, cifar_test_loader, cifar_classes, 'cifar', device)